package org.mockfu.junit;

import static org.mockfu.Exceptions.*;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.lang.reflect.Method;
import java.net.URL;
import java.net.URLDecoder;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

import org.junit.Test;
import org.junit.internal.runners.InitializationError;
import org.junit.runner.RunWith;
import org.junit.runner.Runner;
import org.junit.runner.manipulation.Filter;
import org.junit.runner.manipulation.NoTestsRemainException;
import org.junit.runner.notification.RunNotifier;
import org.junit.runners.Suite;

public class ParallelSuite extends Suite
{

    public ParallelSuite(Class<?> klass) throws InitializationError
    {
        super(klass, getTestClasses(klass));
    }

    @Override
    public void filter(Filter filter) throws NoTestsRemainException
    {
    }

    @Override
    public void run(final RunNotifier notifier)
    {
        List<Runner> runners = tryGetRunners();
        if (runners == null)
        {
            System.err.println("Cannot run tests in parallel: incompatible JUnit version");
            return;
        }


        List<Future<?>> futures = new ArrayList<Future<?>>();

        for (final Runner runner : runners)
        {
            String[] mutexNames = getMutexes(runner);

            if (isParallelized(runner))
            {
                /* a Parallelizer test runner thread can't run in ParallelExecutor or you'll get a
                 * deadlock when pool size is one (its the individual tests, not the runner, that
                 * should run in ParallelExecutor) */
                futures.add(Executors.newSingleThreadExecutor().submit(new Runnable()
                {
                    public void run()
                    {
                        runner.run(notifier);
                    }
                }));
            }
            else
            {
                futures.add(ParallelExecutor.submit(new Runnable()
                {
                    public void run()
                    {
                        runner.run(notifier);
                    }
                }, mutexNames));
            }
        }

        /* wait for each thread to complete */
        try
        {
            for (Future<?> each : futures)
                each.get();
        }
        catch (Exception ex)
        {
            throw undeclared(ex);
        }
    }


    private static String[] getMutexes(Runner runner)
    {
        if (isParallelized(runner))
            /* let parallelizer handle the mutexes */
            return null;

        Set<String> mutexes = new HashSet<String>();

        Serialize serialize = runner.getDescription().getAnnotation(Serialize.class);
        if (serialize != null)
        {
            for (String mutex : serialize.around())
            {
                if (mutex.contains(","))
                    for (String each : mutex.split(","))
                        mutexes.add(each);
                else
                    mutexes.add(mutex);
            }
        }

        List<Method> methods = getTestMethods(runner);
        if (methods != null)
        {
            /* get the union of all method mutexes */
            for (Method method : methods)
            {
                serialize = method.getAnnotation(Serialize.class);
                if (serialize != null)
                {
                    for (String mutex : serialize.around())
                    {
                        if (mutex.contains(","))
                            for (String each : mutex.split(","))
                                mutexes.add(each);
                        else
                            mutexes.add(mutex);
                    }
                }
            }
        }

        return mutexes.toArray(new String[mutexes.size()]);
    }

    private static boolean isParallelized(Runner runner)
    {
        RunWith runWith = runner.getDescription().getAnnotation(RunWith.class);
        return runWith != null && runWith.value() == Parallelizer.class;
    }


    @SuppressWarnings("unchecked")
    private static List<Method> getTestMethods(Runner runner)
    {
        List<Method> methods = null;
        for (Class<?> klass = runner.getClass(); klass != Object.class; klass = klass.getSuperclass())
        {
            try
            {
                /* protected method on JUnitClassRunner */
                Method getMethodsMethod = klass.getDeclaredMethod("getTestMethods");
                getMethodsMethod.setAccessible(true);
                methods = (List<Method>)getMethodsMethod.invoke(runner);
            }
            catch (NoSuchMethodException ex)
            {
                continue;
            }
            catch (Exception ex)
            {
                throw undeclared(ex);
            }
        }
        return methods;
    }

    @SuppressWarnings("unchecked")
    private List<Runner> tryGetRunners()
    {
        try
        {
            return (List<Runner>)this.getClass().getMethod("getRunners").invoke(this);
        }
        catch (NoSuchMethodException ex)
        {
            return null;
        }
        catch (Exception ex)
        {
            throw undeclared(ex);
        }
    }

    private static Class<?>[] getTestClasses(Class<?> klas)
    {
        /* initialize parallel executor thread count (before test classes are loaded) */
        Parallel parallel = klas.getAnnotation(Parallel.class);
        int threads = (parallel == null ? -1 : parallel.threads());
        ParallelExecutor.initialize(threads);

        SuiteClasses annotation= klas.getAnnotation(SuiteClasses.class);
        if (annotation != null)
            return annotation.value();

        List<Class<?>> tests = new ArrayList<Class<?>>();
        ClassLoader loader = klas.getClassLoader();
        String packageName = "";
        if (klas.getPackage() != null)
            packageName = klas.getPackage().getName().replace('.', '/');

        try
        {

            for (Enumeration<URL> urls = loader.getResources(packageName); urls.hasMoreElements(); )
            {
                String url = urls.nextElement().getFile();
                url = URLDecoder.decode(url, "UTF-8");

                if (url.startsWith("file:"))
                    url = url.substring(5);
                else if (url.contains("!"))
                    url = url.substring(0, url.indexOf('!'));

                File file = new File(url);

                if (file.isDirectory())
                {
                    String classFilename = klas.getSimpleName() + ".class";
                    for (String each : file.list())
                    {
                        if (each.equals(classFilename))
                            loadFromFileSystem(tests, packageName, file, loader);
                    }
                }
                else
                    loadFromJar(tests, packageName, file, loader);
            }
        }
        catch (Exception ex)
        {
            throw undeclared(ex);
        }

        tests.remove(klas);

        return tests.toArray(new Class[tests.size()]);
    }

    private static void loadFromFileSystem(List<Class<?>> tests, String path, File directory, ClassLoader loader)
    {
        for (File file : directory.listFiles())
        {
            String resource = file.getName();
            if (!path.equals(""))
                resource = path + "/" + resource;

            if (file.isDirectory())
            {
                loadFromFileSystem(tests, resource, file, loader);
            }
            else if (file.getName().endsWith(".class"))
            {
                addIfTest(tests, resource, loader);
            }
        }
    }


    private static void loadFromJar(List<Class<?>> tests, String path, File jarfile, ClassLoader loader)
    {
        try
        {
            JarInputStream in = new JarInputStream(new FileInputStream(jarfile));
            try
            {
                for (JarEntry entry; (entry = in.getNextJarEntry()) != null; )
                {
                    String name = entry.getName();

                    if (!entry.isDirectory() && name.startsWith(path) && name.endsWith(".class"))
                    {
                        addIfTest(tests, name, loader);
                    }
                }
            }
            finally
            {
                in.close();
            }
        }
        catch (IOException ex)
        {
            throw undeclared(ex);
        }
    }

    /**
     * Adds the class if it has a RunWith annotation or a method with the Test annotation.
     */
    private static void addIfTest(List<Class<?>> tests, String resource, ClassLoader loader)
    {
        try
        {
            Class<?> klas = loader.loadClass(resource.replace('/', '.').replace(".class", ""));
            RunWith runWith = klas.getAnnotation(RunWith.class);
            if (runWith != null)
            {
                /* add unless another ParallelSuite */
                if (runWith.value() != ParallelSuite.class)
                    tests.add(klas);
                return;
            }
            for (Method method : klas.getDeclaredMethods())
            {
                if (method.getAnnotation(Test.class) != null)
                {
                    tests.add(klas);
                    return;
                }
            }
        }
        catch (Exception ex)
        {
            undeclared(ex);
        }
    }

}
